package com.langll.ai.milvus;

import java.util.Collections;

import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import com.langll.ai.AIChatApplication;
import com.langll.ai.vector.config.MilvusArchive;
import com.langll.ai.vector.service.MilvusService;

import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.index.request.CreateIndexReq;
import lombok.extern.slf4j.Slf4j;

/**
 * 操作向量数据库
 */
@Slf4j
@SpringBootTest(classes = AIChatApplication.class)
@RunWith(SpringRunner.class)
public class TestMiluse {


    @Autowired
    MilvusClientV2 milvusClientV2;

    @Test
    public void createCollection(){
// 创建字段
        CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema()
                // 创建主键字段
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.ID)
                        // 字段描述
                        .description("主键ID")
                        // 字段类型
                        .dataType(DataType.Int64)
                        // 是否为主键
                        .isPrimaryKey(true)
                        // 设置主键自增
                        .autoID(true)
                        .build())
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.FILE_NAME)
                        // 字段描述
                        .description("文件名")
                        // 字段类型
                        .dataType(DataType.VarChar)
                        // 设置字段为可空
                        .isNullable(true)
                        .build())
                // 创建特征向量字段
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.FEATURE)
                        // 字段描述
                        .description("特征向量")
                        // 字段类型
                        .dataType(DataType.FloatVector)
                        // 设置向量维度
                        .dimension(MilvusArchive.FEATURE_DIM)
                        .build())
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.TEXT)
                        // 字段描述
                        .description("文本")
                        // 字段类型
                        .dataType(DataType.VarChar)
                        // 设置字段为可空
                        .isNullable(true)
                        .build())
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.METADATA)
                        // 字段描述0
                        .description("元数据")
                        // 字段类型
                        .dataType(DataType.VarChar)
                        // 设置字段为可空
                        .isNullable(true)
                        .build());
        // 创建集合
        CreateCollectionReq collectionReq = CreateCollectionReq.builder()
                // 集合名称
                .collectionName(MilvusArchive.COLLECTION_NAME)
                // 集合描述
                .description("自定义知识库")
                // 集合字段
                .collectionSchema(schema)
                // 分片数量
                .numShards(MilvusArchive.SHARDS_NUM)
                .build();
        milvusClientV2.createCollection(collectionReq);

        // 创建索引
        IndexParam indexParam = IndexParam.builder()
                // 索引字段名
                .fieldName(MilvusArchive.Field.FEATURE)
                // 索引类型
                .indexType(IndexParam.IndexType.IVF_FLAT)
                // 索引距离度量
                .metricType(IndexParam.MetricType.COSINE)
                .build();
        CreateIndexReq createIndexReq = CreateIndexReq.builder()
                .collectionName(MilvusArchive.COLLECTION_NAME)
                .indexParams(Collections.singletonList(indexParam))
                .build();

        milvusClientV2.createIndex(createIndexReq);
    }
}
